import torch
from torch import nn
import torch.nn.functional as F


class Double_conv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(Double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.LeakyReLU(negative_slope=0.2, inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class up(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(up, self).__init__()
        self.conv = Double_conv(in_ch, out_ch)

    def forward(self, x1, x2):
        x1 = F.interpolate(x1, scale_factor=2)
        x = torch.cat([x1, x2], dim=1)
        x = self.conv(x)
        return x

class U_net(nn.Module):
    def __init__(self, n_channels, n_classes, ch):
        super(U_net, self).__init__()
        self.sigmoid = nn.Sigmoid()
        self.inc = Double_conv(n_channels, ch)
        self.down1= nn.Sequential(nn.MaxPool2d(2), Double_conv(ch, ch * 2))
        self.down2= nn.Sequential(nn.MaxPool2d(2), Double_conv(ch * 2, ch * 4))
        self.down3= nn.Sequential(nn.MaxPool2d(2), Double_conv(ch * 4, ch * 8))
        self.down4= nn.Sequential(nn.MaxPool2d(2), Double_conv(ch * 8, ch * 8))

        self.up1 = up(ch * 16, ch * 4)
        self.up2 = up(ch * 8, ch * 2)
        self.up3 = up(ch * 4, ch)
        self.up4 = up(ch * 2, ch)
        self.outc = nn.Conv2d(ch, n_classes, 1)

    def forward(self, x1, x2):
        x = torch.cat([x1, x2], 1)
        x1 = self.inc(x)  # [64, 960, 960]
        x2 = self.down1(x1) #[128, 480, 480]
        x3 = self.down2(x2)  # [256, 240, 240]
        x4 = self.down3(x3)  # [512, 120, 120]
        x5 = self.down4(x4)  # [512, 60, 60]
        x = self.up1(x5, x4)  # [512, 120, 120]
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        x = self.outc(x)
        return self.sigmoid(x)


if __name__ == '__main__':
    t1 = torch.ones([1, 3, 960, 960]).cuda()
    t2 = torch.ones([1, 3, 960, 960]).cuda()
    net = U_net(6, 2, 8).cuda()
    print(net(t1, t2).shape)